package com.zjuee.gateway;

import com.zjuee.gateway.utils.SecurityUtil;
import org.apache.shiro.codec.Base64;
import org.apache.shiro.mgt.AbstractRememberMeManager;
import org.apache.shiro.subject.Subject;
import org.apache.shiro.subject.SubjectContext;
import org.apache.shiro.web.servlet.ShiroHttpServletRequest;
import org.apache.shiro.web.subject.WebSubjectContext;
import org.apache.shiro.web.util.WebUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.servlet.ServletRequest;
import javax.servlet.http.HttpServletRequest;

/**
 * 定义RememberMe控制器
 *
 * @author wu liang
 * @since 2020/8/13 9:54
 */
public class MyRememberMe extends AbstractRememberMeManager {
    private static transient final Logger log = LoggerFactory.getLogger(MyRememberMe.class);

    //清除序列化
    @Override
    protected void forgetIdentity(Subject subject) {
    }

    /**
     * 实现序列化，因为是前后端分离，我们将信息记录在header中，
     * 所以这一步我们转为字符串后放入request域，在onLoginSuccess方法中响应到客户端
     */
    @Override
    protected void rememberSerializedIdentity(Subject subject, byte[] serialized) {
        if (!WebUtils.isHttp(subject)) {
            return;
        }
        HttpServletRequest request = WebUtils.getHttpRequest(subject);
        String base64 = Base64.encodeToString(serialized);
        request.setAttribute(SecurityUtil.REMEMBER_REQUEST_KEY, base64);
    }

    /**
     * 读取序列化
     */
    @Override
    protected byte[] getRememberedSerializedIdentity(SubjectContext subjectContext) {
        if (!WebUtils.isHttp(subjectContext)) {
            return null;
        }

        WebSubjectContext wsc = (WebSubjectContext) subjectContext;
        if (isIdentityRemoved(wsc)) {
            return null;
        }
        HttpServletRequest request = WebUtils.getHttpRequest(wsc);

        // 从header中获取rememberMeValue
        String header = request.getHeader(SecurityUtil.REMEMBER_KEY);

        if (header != null) {
            header = ensurePadding(header);
            return Base64.decode(header);
        }

        return null;
    }

    private boolean isIdentityRemoved(WebSubjectContext subjectContext) {
        ServletRequest request = subjectContext.resolveServletRequest();
        if (request != null) {
            Boolean removed = (Boolean) request.getAttribute(ShiroHttpServletRequest.IDENTITY_REMOVED_KEY);
            return removed != null && removed;
        }
        return false;
    }

    @Override
    public void forgetIdentity(SubjectContext subjectContext) {

    }

    private String ensurePadding(String base64) {
        int length = base64.length();
        if (length % 4 != 0) {
            StringBuilder sb = new StringBuilder(base64);
            for (int i = 0; i < length % 4; ++i) {
                sb.append('=');
            }
            base64 = sb.toString();
        }
        return base64;
    }
}
